package com.mjlf.algorithms.search;

import org.junit.Test;

import java.util.ArrayList;
import java.util.List;

/**
 * @auther a123
 * @create 2018-09-12 10:51
 * @desc 二叉查找树
 * 二叉查找树，要求满足左子树小于右子树，其实也快速排序相似，做到左边元素总是小于右边元素
 */
public class BST<Key extends Comparable<Key>, Value> {

    private Node root;

    private class Node {
        private Key key;
        private Value value;
        private Node left, right;
        private int num;

        public Node(Key key, Value value, int num) {
            this.key = key;
            this.value = value;
            this.num = num;
        }

        @Override
        public String toString() {
            return "Node{" +
                    "key=" + key +
                    ", value=" + value +
                    ", left=" + left +
                    ", right=" + right +
                    ", num=" + num +
                    '}';
        }
    }

    public int size() {
        return size(root);
    }

    private int size(Node node) {
        if (node == null) {
            return 0;
        }
        return node.num;
    }

    public Value get(Key key) {
        return get(root, key);
    }

    private Value get(Node node, Key key) {
        if (node == null) {
            return null;
        }
        if (key.compareTo(node.key) == 0) {
            return node.value;
        } else if (key.compareTo(node.key) < 0) {
            return get(node.left, key);
        } else {
            return get(node.right, key);
        }
    }

    public void put(Key key, Value value) {
        root = put(root, key, value);
    }

    private Node put(Node node, Key key, Value value) {
        if (node == null) {
            return new Node(key, value, 1);
        }
        int cmp = key.compareTo(node.key);
        if (cmp < 0) {
            node.left = put(node.left, key, value);
        } else if (cmp > 0) {
            node.right = put(node.right, key, value);
        } else {
            node.value = value;
        }
        node.num = size(node.left) + size(node.right) + 1;
        return node;
    }

    /**
     * 获取包含在左右键之间的所有键
     *
     * @param leftKey
     * @param rightKey
     *
     * @return
     */
    public List<Key> keys(Key leftKey, Key rightKey) {
        List<Key> keys = new ArrayList<>();
        keys(root, leftKey, rightKey, keys);
        return keys;
    }

    private void keys(Node node, Key leftKey, Key rightKey, List<Key> keys) {
        if (node == null) {
            return;
        }
        if (node.key.compareTo(leftKey) > 0) {
            keys(node.left, leftKey, rightKey, keys);
        }
        if (node.key.compareTo(leftKey) >= 0 && node.key.compareTo(rightKey) <= 0) {
            keys.add(node.key);
        }
        if (node.key.compareTo(rightKey) < 0) {
            keys(node.right, leftKey, rightKey, keys);
        }
    }

    public void deleteKey(Key key) {
        root = deleteKey(root, key);
    }

    /**
     * 删除指定元素
     * @param node
     * @param key
     * @return
     */
    private Node deleteKey(Node node, Key key) {
        if(node == null){
            return null;
        }
        int cmp = node.key.compareTo(key);
        if (cmp == 0) {//找到需要删除的元素
            if (node.left == null) {
                node = node.right;
            } else if (node.right == null) {
                node = node.left;
            } else {

                Node right = node.right;
                node = node.left;
                Node lastRight = node;
                while (lastRight.right != null) {
                    lastRight = node.right;
                }
                lastRight.right = right;
            }
        } else if (cmp < 0) {//说明要删除的元素在右子树里
            node.right = deleteKey(node.right, key);
        } else if (cmp > 0) {//说明要删除的元素在左子树里
            node.left = deleteKey(node.left, key);
        }
        return node;
    }

    public void show() {
        show(root);
        System.out.println();
    }

    private void show(Node node) {
        if (node != null) {
            System.out.print(node.key + " ");
            if (node.left != null) {
                show(node.left);
            }
            if (node.right != null) {
                show(node.right);
            }
        }
    }

    @Test
    public void test() {
        BST<Integer, Integer> bst = new BST<>();
        bst.put(5, 5);
        bst.put(3, 5);
        bst.put(2, 5);
        bst.put(1, 5);
        bst.put(4, 5);
        bst.put(10, 5);
        bst.put(7, 5);
        bst.put(6, 5);
        bst.put(8, 5);
        bst.put(9, 5);
        bst.put(15, 5);
        bst.put(12, 5);

        bst.show();
        bst.deleteKey(3);
        bst.show();
    }
}
